import cv2
import numpy as np
import kociemba
import colorsys

# Define the colors for comparison
color_ranges = [
    # c    h1     h2         s1    s2 invert
    ['O',  .04,  0.12, .4,  2, False],  # Orange
    ['R', .9, .02, .4,  2, True],  # Red
    ['G',  .33,  .45, .7,  2, False],  # Green
    ['Y',  .12, .24, .5,  .8, False],  # Yellow
    ['B',  .55,  .61, .6,  2, False],  # Blue
    ['W',  -1,  1,  -1, .4, False]  # White
]

# Define the coordinates of the 9 squares on the Rubik's Cube
squares = {
    1: [(100, 100), (200, 200)],
    2: [(225, 100), (325, 200)],
    3: [(350, 100), (450, 200)],
    4: [(100, 225), (200, 325)],
    5: [(225, 225), (325, 325)],
    6: [(350, 225), (450, 325)],
    7: [(100, 350), (200, 450)],
    8: [(225, 350), (325, 450)],
    9: [(350, 350), (450, 450)]
}

# Initialize the Rubik's Cube faces
rubiks_cube = {
    'U': [''] * 9,
    'F': [''] * 9,
    'R': [''] * 9,
    'B': [''] * 9,
    'L': [''] * 9,
    'D': [''] * 9
}

# Initialize the current face index
current_face_index = 0
faces = list(rubiks_cube.keys())
current_face = faces[current_face_index]


# Callback function for key press event
def key_press(event):
    global current_face_index, current_face

    if event == ord('q'):
        cv2.destroyAllWindows()
    elif event == ord(' '):
        current_face_index += 1
        if current_face_index >= len(faces):
            current_face_index = 0
        current_face = faces[current_face_index]


# Capture video from webcam
cap = cv2.VideoCapture(3)
print('e')
# Set the resolution of the video
cap.set(3, 640)
cap.set(4, 480)

# Flag to track if all faces are scanned
all_faces_scanned = False

while True:
    # Read the current frame from the video capture
    ret, frame = cap.read()

    # Draw the Rubik's Cube squares overlay on the frame
    for square, coords in squares.items():
        x1, y1 = coords[0]
        x2, y2 = coords[1]

        # Extract the region of interest (square)
        roi = frame[y1:y2, x1:x2]

        # Calculate the average color of the square
        avg_color = np.median(roi, axis=(0, 1))
        # Compare the average color with the predefined colors
        hsv_color = colorsys.rgb_to_hsv(avg_color[2]/255, avg_color[1]/255, avg_color[0]/255)
        print(hsv_color[0])
        closest_color = "aaa"
        for color in color_ranges:
            if not color[5]:
                if color[1] < hsv_color[0] < color[2] and color[3] < hsv_color[1] < color[4]:
                    closest_color = color[0]
                    break
            else:
                if (hsv_color[0] > color[1] or hsv_color[0] < color[2]) and color[3] < hsv_color[1] < color[4]:
                    closest_color = color[0]
                    break


        # color_diffs = [np.linalg.norm(avg_color - np.array(color)) for color in colors.values()]
        # closest_color = list(colors.keys())[color_diffs.index(min(color_diffs))]

        # Draw the square overlay
        cv2.rectangle(frame, coords[0], coords[1], (0, 255, 0), 2)
        cv2.putText(frame, closest_color, (int((x1 + x2) / 2 - 10), int((y1 + y2) / 2 + 10)),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(frame, "(" + str(round(hsv_color[0] * 1000)) + ", " + str(round(hsv_color[1] * 100)) + ")",
                    (int((x1 + x2) / 2 - 10), int((y1 + y2) / 2 + 40)),
                    cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 255, 0), 1)

        # Store the color for the corresponding square on the current face
        rubiks_cube[current_face][square - 1] = closest_color

    # Display the current face on the frame
    cv2.putText(frame, 'Scanning Face: ' + current_face, (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    # Show the frame
    cv2.imshow('Rubik\'s Cube Scanner', frame)

    # Wait for key press event
    key = cv2.waitKey(1) & 0xFF

    # Process key press event
    if key != 255:
        key_press(key)

    # Check if 'q' key is pressed to quit
    if key == ord('q'):
        break

    # Check if all faces are scanned
    if all(rubiks_cube[face] != [''] * 9 for face in rubiks_cube):
        all_faces_scanned = True

    # Break the loop if all faces are scanned and current face is 'U'
    if all_faces_scanned and current_face == 'U':
        break

# Release the video capture and close windows
cap.release()
cv2.destroyAllWindows()

# Combine the face colors to form the cube state
cube_state = [rubiks_cube['U'], rubiks_cube['R'], rubiks_cube['F'],
              rubiks_cube['D'], rubiks_cube['L'], rubiks_cube['B']]

# Map the color changes
color_mapping = {
    rubiks_cube['U'][4]: 'U',
    rubiks_cube['L'][4]: 'L',
    rubiks_cube['F'][4]: 'F',
    rubiks_cube['R'][4]: 'R',
    rubiks_cube['B'][4]: 'B',
    rubiks_cube['D'][4]: 'D'
}

# Convert the cube state to a single line of 54 letters without spaces
state_line = ''.join([color_mapping[color] for face in cube_state for color in face])
# Print the cube state
print(state_line)
# Define the initial state of the Rubik's Cube using the colors/pieces
initial_state = state_line

# Solve the cube using the kociemba package
solution = kociemba.solve(initial_state)
# solution = "R2 D2 F2 D2 B3 R2 B1 D3 R2 F1 L3 B2 F3 D3 L1 F3 R3 F2 D3 (19f)"
# Reformat the string into a list
final_solution = []

i = 0
while True:
    # woo if statement spam
    if solution[i] == "1":
        final_solution[-1] += "n"
        i += 2
    elif solution[i] == "3":
        final_solution[-1] += "p"
        i += 2
    elif solution[i] == "2":
        final_solution[-1] += "n"
        final_solution.append(final_solution[-1])
        i += 2
    elif solution[i] == "(":
        break
    else:
        final_solution.append(solution[i])
        i += 1
# check for edge case if final move was a normal move
if len(final_solution[-1]) == 1:
    final_solution[-1] += "n"

# Print the original solution and the final solution
print("Original Solution:", solution)
print("Modified Solution: ", final_solution)